DP还能这么跑?基于非冯·诺依曼架构的DP高速计算
长期以来,在“速度”和“精度”这两个核心指标上,分子动力学(molecular dynamics,MD)存在“鱼与熊掌不可兼得”的问题——经典分子动力学(classical molecular dynamics, CMD)速度快,但精度低,难以满足高精度计算要求;第一性原理分子动力学(ab-initio molecular dynamics, AIMD)精度高,但速度慢,难以计算大系统。
近年来,DP团队提出了基于深度学习势的分子动力学(DeePMD),在保持AIMD级别精度前提下,大幅提升了计算速度。然而,DeePMD的计算速度虽然比AIMD快多个数量级,但仍比CMD慢2个数量级左右。研发出兼具高精度、高速度的MD方法,一直是DP社区(乃至整个MD学术界)追求的目标。
近日,湖南大学刘杰教授课题组提出了基于新型非冯·诺依曼架构的分子动力学(non von-Neumann molecular dynamics, NVNMD),同时实现了AIMD级别的高精度、CMD级别的高速度——实现了“鱼与熊掌可以兼得”。该成果以“Accurate and efficient molecular dynamics based on machine learning and non von Neumann architecture”发表在了npj Computational Materials(第一作者为莫平辉,通讯作者为刘杰)¹,训练代码和计算服务均已公开发布,目前免费提供给大家使用²’³。
下面,我们一起来看看,“鱼与熊掌兼得”是怎么实现的吧。
01
问题根源——冯·诺依曼架构和“存储墙”瓶颈
自1946年发明至今,冯·诺依曼架构一直占据统治地位,是手机、台式机、笔记本、计算服务器、超级计算中心的底层基础架构。目前,需要运行MD计算(包括DeePMD计算)时,使用冯·诺依曼(von Neumann, vN)架构计算机是几乎所有研究人员的唯一选择,这已成为一种“固有范式”(paradigm)。
遗憾的是,冯·诺依曼架构中,计算单元(例如CPU/GPU)和存储单元(例如内存)是互相独立的(即“存算分离”),导致计算总耗时和总功耗的绝大部分(>90%)消耗于存储单元、计算单元之间的频繁数据搬运,俗称“存储墙(memory wall)”和“功耗墙(power wall)”瓶颈。这严重制约了MD计算性能的提升。
02
解决方案——新型非冯·诺依曼架构的NVNMD
为解决该问题,湖南大学研究团队自主研制了基于新型非冯·诺依曼芯片架构的MD计算系统“NVNMD”。NVNMD的核心计算模块中,存储单元和计算单元紧密融为一体(即“存算一体”),避免了频繁的数据搬运,极大缓解了MD计算中的“存储墙”和“功耗墙”瓶颈。因此,相较部署于GPU的DeePMD,将计算速度提升了约1-2个数量级;将计算功耗降低了2-3个数量级。
在算法层面,为实现从传统CPU/GPU向新型非冯·诺依曼芯片架构的“范式转移 (paradigm shift) ”,湖南大学研究团队对DeePMD进行了若干修改和创新:(1)用整数等离散数据类型,取代了浮点数等连续数据类型;(2)用移位等运算,取代了乘法等运算;(3)用离散神经网络,取代了连续神经网络;(4)用逼近近似,取代了三角函数求解;等。这些修改和创新旨在使用新型非冯·诺依曼架构下有限的硬件资源,实现尽可能高的片内并行度,以实现高速MD计算。
03
技术细节——NVNMD系统架构
每一步MD计算包含两部分:第①部分,计算高维势能面,包括能量、受力以及维里的预测;第②部分,其他操作,包括构建相邻列表,牛顿力学积分等。第①部分计算功能单一、计算量大;第②部分计算功能复杂、计算量小。因此,在功能划分上,速度快但不灵活的SPU用来执行第①部分计算;速度慢但计算灵活的MPU用来执行第②部分计算;HTI则被用来保证MPU和SPU之间高速通信。
MPU中的多个计算核心以流水线方式调用SPU进行DP模型推理,使得SPU计算效率最大化。
SPU利用存内计算和流水线计算技术,使得DP模型保存在片内存储器,免除在每步MD读取模型数据;并且前一步的计算结果直接作为下一步计算的输入,消除了处理器与片外存储器之间的频繁的数据传输,大大加快计算速度,降低计算能耗。
HTI保证MPU和SPU之间的高带宽数据传输,使得两者的通信不会成为新的性能瓶颈。
04
性能分析——NVNMD的速度、精度、功耗
为了检验NVNMD方法的性能,湖南大学研究团队使用CPU(Intel i7-10700K)和FPGA(Xilinx xcvu9p)硬件来实现图 1中描述的计算机系统。CPU和FPGA间使用PCIe 3.0×16作为通信接口。在训练阶段,DP模型使用CPU/GPU硬件进行训练。在推理阶段,CPU运行LAMMPS软件进行分子动力学仿真,并通过自定义的用户接口调用FPGA进行势能面的推理。为全面评估实例性能,从精度,速度以及能效3个方面,进行了定量测试。
▲ 图2 · 三种MD方法在精度和速度上的对比
如图 3a所示,使用NVNMD运行GeTe的分子动力学,复现了GeTe的整个相变过程(晶体相→液相→无定形相→晶体相)⁶。如图 4所示,其无定形相的径向分布函数,角度分布函数以及配位数与DeePMD和AIMD结果基本一致。
如图 3b所示,使用NVNMD运行Li₁₀Ge₂SP₁₂的分子动力学,得到的500K下的Li原子的均方位移(mean square displacement,MSD)为2.03×10⁻¹⁰m²s⁻¹,与文献值⁷(约2.00×10⁻¹⁰m²s⁻¹)基本一致。
▲ 图3 · (a) GeTe相变过程 (b) Li₁₀Ge₂PS₁₂中Li原子的均方位移
▲ 图4 · 无定形相GeTe的MD轨迹结构性质对比,包括径向分布函数(左),角度分布函数(中),和配位数(右)
利用非冯·诺依曼架构的优势,NVNMD不仅可以大幅提升计算速度,还可以大幅降低计算能耗。如表 1所示,NVNMD在保持高速度、高精度的同时,可将能效提升2-3个数量级。
05
更多资源——代码和论文链接
想尝试NVNMD的朋友们,欢迎点击原文链接进行报名,因为资源有限,我们会尽快向您提供测试账号。同时您可以通过以下链接下载训练代码(https://github.com/LiuGroupHNU/nvnmd)来训练模型,之后通过提供的账号登录服务器(http://nvnmd.picp.vip/),上传DP模型和lammps输入脚本来进行测试。上述资源都是免费开放的。NVNMD的训练代码会在近期合并到DeePMD的GitHub账号。
希望了解NVNMD更多细节的朋友,欢迎点击论文(全文可免费下载)链接:https://www.nature.com/articles/s41524-022-00773-z
参考文献
[1] Mo, P. et al. Accurate and efficientmolecular dynamics based on machine learning and non von Neumann architecture. NpjComput. Mater. 8, 1–15 (2022).
[2] Liu, J. & Mo, P. The training andtesting code for NVNMD. https://github.com/LiuGroupHNU/nvnmd (2021).
[3] Liu, J. & Mo, P. The server website ofNVNMD. http://nvnmd.picp.vip/ (2021).
[4] Zhang, L. et al. End-to-end SymmetryPreserving Inter-atomic Potential Energy Model for Finite and Extended Systems.Adv. Neural Inf. Process. Syst. 2018-Decem, 4436–4446 (2018).
[5] Wang, H., Zhang, L., Han, J. & E, W.DeePMD-kit: A deep learning package for many-body potential energyrepresentation and molecular dynamics. Comput. Phys. Commun. 228,178–184 (2018).
[6] Shi, M., Li, J., Tao, M., Zhang, X. &Liu, J. Artificial intelligence model for efficient simulation of monatomicphase change material antimony. Mater. Sci. Semicond. Process. 136,106146 (2021).
[7] Huang, J. et al. Deep potentialgeneration scheme and simulation protocol for the Li 10 GeP 2 S 12 -typesuperionic conductors. J. Chem. Phys. 154, 094703 (2021).
[8] Jia, W. et al. Pushing the Limit ofMolecular Dynamics with Ab Initio Accuracy to 100 Million Atoms with MachineLearning. in SC20: International Conference for High Performance Computing,Networking, Storage and Analysis 1–14 (IEEE, 2020).doi:10.1109/SC41405.2020.00009.
- End -
DP还能这么跑
DP用于TCAD中构建半导体掺杂原子迁移模型
DP还能干这个?深度势能描述氘的超固体相
DP还能干这个?DP模拟揭示盐溶液的微观结构